package cn.powerchaintech.common.tcp;

import cn.powerchaintech.common.tcp.exception.MessageException;
import com.google.protobuf.InvalidProtocolBufferException;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.zeromq.SocketType;
import org.zeromq.ZContext;
import org.zeromq.ZMQ;

import java.io.Closeable;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * 服务端
 */
public class TcpServer implements Closeable {
    private static final Logger logger = LogManager.getLogger(TcpServer.class);
    private final int routerPort;
    private final int pubPort;
    private final Map<String, byte[]> identities = Collections.synchronizedMap(new HashMap<>());
    private ZContext context;
    private ZMQ.Socket routerSocket;
    private ZMQ.Socket pubSocket;
    public ResponseCallback cb;

    public TcpServer(int port) {
        this.routerPort = port;
        this.pubPort = port + 1;
    }

    public void start() {
        context = new ZContext();
        routerSocket = context.createSocket(SocketType.ROUTER);
        routerSocket.bind("tcp://*:" + routerPort);

        // 创建发布-订阅模式的套接字
        pubSocket = context.createSocket(SocketType.PUB);
        pubSocket.bind("tcp://*:" + pubPort);

        ZMQ.Poller poller = context.createPoller(1);
        poller.register(routerSocket, ZMQ.Poller.POLLIN);
        ExecutorService executor = Executors.newFixedThreadPool(1);
        executor.execute(() -> {
            while (!Thread.currentThread().isInterrupted()) {
                poller.poll();
                if (poller.pollin(0)) {
                    // 处理路由模式的消息
                    boolean isNewClient = true;
                    byte[] identity = routerSocket.recv();
                    for (byte[] id : identities.values()) {
                        if (Arrays.equals(id, identity)) {
                            isNewClient = false;
                            break;
                        }
                    }
                    if (isNewClient) {
                        String appName = routerSocket.recvStr();
                        identities.put(appName, identity);
                        continue;
                    }
                    byte[] data = routerSocket.recv();
                    try {
                        Message.Messenger messenger = Message.Messenger.parseFrom(data);
                        if(cb != null) {
                            cb.recv("-", messenger.getBody().toByteArray());
                        }
                    } catch (InvalidProtocolBufferException ex) {
                        throw new MessageException(ex);
                    }
                }
            }
        });
        executor.shutdown();
    }

    public void send(String appName, byte[] data) {
        byte[] identity = identities.get(appName);
        if(identity != null) {
            routerSocket.sendMore(identity);
            routerSocket.send(data);
        }
    }

    @Override
    public void close() throws IOException {
        routerSocket.close();
        pubSocket.close();
        context.close();
    }
}
